package com.ruoyi.ade.service.impl;

import com.alibaba.fastjson2.JSONArray;
import com.ruoyi.ade.domain.*;
import com.ruoyi.ade.llm.model.Ade;
import com.ruoyi.ade.llm.model.AdeRetrivedKnowledge;
import com.ruoyi.ade.llm.service.IAdeKnowledgeEmbedService;
import com.ruoyi.ade.llm.service.IAdeService;
import com.ruoyi.ade.mapper.AdePromptMapper;
import com.ruoyi.ade.mapper.AdeRagLabMapper;
import com.ruoyi.ade.mapper.AdeRagLabResultMapper;
import com.ruoyi.ade.mapper.StandardMapper;
import com.ruoyi.ade.model.AdeKnowledgeRetriveRequest;
import com.ruoyi.ade.service.IAdeKnowledgeCacheService;
import com.ruoyi.ade.service.IAdeRagExecuteService;
import com.ruoyi.ade.service.IAdeRagLabService;
import com.ruoyi.ade.tool.AdeCompareUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
@Slf4j
public class AdeRagExecuteServiceImpl implements IAdeRagExecuteService {

    @Autowired
    private AdeRagLabMapper adeRagLabMapper;

    @Autowired
    private AdePromptMapper adePromptMapper;

    @Autowired
    private StandardMapper standardMapper;

    @Autowired
    private AdeRagLabResultMapper adeRagLabResultMapper;

    @Autowired
    private IAdeKnowledgeEmbedService adeKnowledgeEmbedService;

    @Autowired
    private IAdeKnowledgeCacheService adeKnowledgeCacheService;

    @Autowired
    private IAdeService adeService;

    @Override
    public void execute(Long id) {

        //验证实验是否存在
        AdeRagLab adeRagLab = adeRagLabMapper.selectAdeRagLabById(id);
        if (adeRagLab == null) throw new RuntimeException("No rag lab found for:" + id);

        //验证实验状态
        if (!AdeLab.STATUS_READY.equals(adeRagLab.getStatus()) && !AdeLab.STATUS_FAILED.equals(adeRagLab.getStatus()))
            throw new RuntimeException("Current lab status exception:" + adeRagLab.getStatus());

        //检查prompt是否存在
        AdePrompt adePrompt = adePromptMapper.selectAdePromptById(adeRagLab.getPromptId());
        if (adePrompt == null) throw new RuntimeException("No prompt found for:" + adeRagLab.getPromptId());

        //更新为运行状态
        updateRagLabStatus(id, AdeLab.STATUS_RUNNING);
        log.info("RagLab id:{}, Started update status:running", id);

        //开始执行任务 线程
        new Thread(new Runnable() {
            @Override
            public void run() {
                long hit_standard = 0;
                try {
                    //Step one: 大模型分析标准数据集
                    log.info("RagLab id:{}, Started step one: LLM analysis standard data", id);
                    List<Standard> standards = standardMapper.selectTopStandardList(adeRagLab.getTotal());
                    if (standards.size() != adeRagLab.getTotal()) {
                        updateRagLabStatus(id, AdeLab.STATUS_FAILED);
                        log.error("RagLab id:{}, Not enough standard data, total:{}, current:{}", id, adeRagLab.getTotal(), standards.size());
                        return;
                    }

                    //获取已经分析完成的记录
                    AdeRagLabResult cond = new AdeRagLabResult();
                    cond.setRagLabId(id);
                    List<AdeRagLabResult> results = adeRagLabResultMapper.selectAdeRagLabResultList(cond);
                    Map<Long, Byte> cache = new HashMap<>();
                    results.forEach(r -> cache.put(r.getStandardId(), (byte) 01));
                    log.info("RagLab id:{}, Get finished results size:{}", id, cache.size());

                    //逐个分析标准数据集
                    String userPrompt, systemPrompt;
                    int finished = 0;

                    for (Standard standard : standards) {
                        if (cache.get(standard.getId()) != null) continue;  //已经分析过了。


                        log.info("RAGLab id:{}, Start analyze standard:{}", id, standard.getId());
                        List<Ade> references = JSONArray.parseArray(standard.getAdes(), Ade.class);   //标准的标注结果  用于后面比对

                        //记录开始时间
                        long cur = System.currentTimeMillis();
                        //每条分析结果
                        AdeRagLabResult adeRagLabResult = new AdeRagLabResult();
                        adeRagLabResult.setStandardId(standard.getId());
                        adeRagLabResult.setRagLabId(id);
                        adeRagLabResult.setStandardAdes(standard.getAdes());
                        String remark = "";


                        float avg_simscore = 0f;
                        int total_knowledges = 0;
                        //构建用户提示
                        userPrompt = String.format(adePrompt.getUserPrompt(), standard.getContent());
                        systemPrompt = adePrompt.getSystemPrompt();

                        //构建rag请求配合
                        AdeKnowledgeRetriveRequest ragRequest = new AdeKnowledgeRetriveRequest();
                        ragRequest.setAdekbid(adeRagLab.getAdekbId());
                        ragRequest.setChunkType(adeRagLab.getChunkType());
                        ragRequest.setSimScoreThreshold(adeRagLab.getSimScore());
                        ragRequest.setTopN(adeRagLab.getTopn().intValue());

                        //检索知识
                        List<AdeRetrivedKnowledge> knowledges = new ArrayList<>();
                        try {
                            //知识库检索
                            List<AdeRetrivedKnowledge> tmps = adeKnowledgeCacheService.retrieveKnowledge(standard.getId(), standard.getContent(), ragRequest);
                            for(AdeRetrivedKnowledge tmp: tmps){
                                if(tmp.getSimScore()>=adeRagLab.getSimScore()) knowledges.add(tmp);
                            }
                            for (AdeRetrivedKnowledge k : knowledges) avg_simscore += k.getSimScore();
                            total_knowledges = knowledges.size();
                            if (total_knowledges != 0)
                                adeRagLabResult.setAvgSimscore(avg_simscore / total_knowledges);
                            else
                                adeRagLabResult.setAvgSimscore(0f);
                            adeRagLabResult.setKnowledgeCount((long) total_knowledges);
                            if(total_knowledges>0) hit_standard++;  //统计检索到知识的标准数据集数量
                        } catch (Exception e) {
                            e.printStackTrace();
                            log.error("RAGLab id:{}, Exception:{}", id, e.getMessage());
                            adeRagLabResult.setAdes("[]");
                            adeRagLabResult.setRemark(e.getMessage());
                            adeRagLabResult.setMatchLevel(AdeCompareUtils.LEVEL_LN);
                            adeRagLabResult.setStatus(AdeLlmLabResult.STATUS_ERR);
                            log.error("RAGLab id:{}, Exception:{}", id, e.getMessage());
                            finished++;
                            adeRagLabResultMapper.insertAdeRagLabResult(adeRagLabResult);
                            updateRagLabProgress(id, (long) finished,hit_standard);
                            continue;
                        }

                        //构建大模型请求
                        systemPrompt = String.format(adePrompt.getSystemPrompt(), AdeRetrivedKnowledge.formatKnowledge(knowledges));
                        //执行ADE识别
                        try {
                            List<Ade> ades = adeService.recognize(adeRagLab.getLlm(), systemPrompt, userPrompt, true);
                            String level = AdeCompareUtils.compare(references, ades);
                            adeRagLabResult.setMatchLevel(level);
                            adeRagLabResult.setAdes(JSONArray.toJSONString(ades));
                            adeRagLabResult.setStatus(AdeLlmLabResult.STATUS_OK);
                            adeRagLabResult.setTimeConsumed((System.currentTimeMillis() - cur) / 1000);
                            finished++;
                            adeRagLabResultMapper.insertAdeRagLabResult(adeRagLabResult);
                            updateRagLabProgress(id, (long) finished,hit_standard);
                        } catch (Exception e) {
                            e.printStackTrace();
                            log.error("RAGLab id:{}, Exception:{}", id, e.getMessage());
                            adeRagLabResult.setAdes("[]");
                            adeRagLabResult.setRemark(e.getMessage());
                            adeRagLabResult.setMatchLevel(AdeCompareUtils.LEVEL_LN);
                            adeRagLabResult.setStatus(AdeLlmLabResult.STATUS_ERR);
                            finished++;
                            adeRagLabResultMapper.insertAdeRagLabResult(adeRagLabResult);
                            updateRagLabProgress(id, (long) finished,hit_standard);
                        }
                    }
                }catch (Exception e){
                    e.printStackTrace();
                    updateRagLabStatus(id, AdeLab.STATUS_FAILED);
                    log.error("RagLab id:{}, Exception:{}", id, e.getMessage());
                }

                //统计分析结果
                AdeRagLabResult adeRagLabResult = new AdeRagLabResult();
                adeRagLabResult.setRagLabId(id);
                List<AdeRagLabResult> adeRagLabResults = adeRagLabResultMapper.selectAdeRagLabResultList(adeRagLabResult);
                int tp =0, tn=0 ,fp=0 ,fn=0;
                long total_knowledge = 0;
                float total_simscore = 0;
                long pos=0,neg=0;
                for(AdeRagLabResult ragLabResult: adeRagLabResults){
                    total_knowledge += ragLabResult.getKnowledgeCount();
                    total_simscore += (ragLabResult.getAvgSimscore()*ragLabResult.getKnowledgeCount());
                    if("[]".equals(ragLabResult.getStandardAdes())){   //negative
                        neg++;
                        if(AdeCompareUtils.LEVEL_LN.equals(ragLabResult.getMatchLevel())){   //没识别出来
                             fp++;
                        }else{
                            tn++;
                        }
                    }else{   //positive
                        pos++;
                        if(AdeCompareUtils.LEVEL_LN.equals(ragLabResult.getMatchLevel())){  //没识别出来
                            fn++;
                        }else{
                            tp++;
                        }
                    }
                }
                float precision = (float)tp/(tp+fp);
                float recall = (float)tp/(tp+fn);
                float f1 = 2*precision*recall/(precision+recall);
                AdeRagLab update = new AdeRagLab();
                update.setId(id);
                update.setPrecisionMean(precision);
                update.setPrecisonSd(0f);
                update.setRecallMean(recall);
                update.setRecallSd(0f);
                update.setF1Mean(f1);
                update.setF1Sd(0f);
                update.setPos(pos);
                update.setNeg(neg);
                update.setHitSimscoreTotal(total_simscore);
                update.setHitKnowledgeCount(total_knowledge);
                update.setStatus(AdeLab.STATUS_SUCCESS);
                adeRagLabMapper.updateAdeRagLab( update);
                log.info("RagLab id:{}, Finished update status:success", id);
            }
        }).start();
    }

    private void updateRagLabProgress(Long id, Long finished, Long hit_standard){
        AdeRagLab adeRagLab = new AdeRagLab();
        adeRagLab.setId(id);
        adeRagLab.setFinished(finished);
        if(hit_standard!=null) adeRagLab.setHitStandard(hit_standard);
        adeRagLabMapper.updateAdeRagLab(adeRagLab);
    }

    private void updateRagLabStatus(Long id, String status) {
        AdeRagLab adeRagLab = new AdeRagLab();
        adeRagLab.setId(id);
        adeRagLab.setStatus(status);
        adeRagLabMapper.updateAdeRagLab(adeRagLab);
    }
}


